There was a great Washington Post article on social distancing. It comes with simulations to show you the effect of different measures, and I think it's fantastic #SciCom. I want to talk it over with my kids, but I want to be ready for them to ask questions like "what if we did (some cool other idea) for social distancing?" So, before I talk it over with them, I want to have some Python code to replicate it and play with. Here are my initial thoughts:

When I tweeted about it, I thought it looked like hard sphere MD. And maybe there are some good stat mech results to be had from that. But I think it looks a bit more like agent based modeling.
I want this to be as reproducible as possible, so I'm tempted to use just the scientific python stack. I'm also tempted to use mesa which looks awesome.
Ugh. I think this means I write it as an object oriented setup, which makes it a bit less accessible.
I'm going to make it a bunch of nested loops, which will be inefficient, but will more easily support different agent types.
I'll write one notebook with all of the code, and then another that has it wrapped up as a package.
I'll also wrap it up in streamlit because OMFG
- Note to self: if I do use streamlit, it looks like it might be smart to just stuff everything into a pandas dataframe. Here's an example of wrapping up a matplotlib animation in streamlit.
Since physics, model the non-movers as just very high mass
NOTE TO SELF: run is fast enough for like 100-1000 particles. The animation is slow. That means I want to pregenerate the trajectory, and then graph that later.

Here's the python setup, using anaconda

conda create --name covid python=3 jupyter jupyterlab numpy scipy matplotlib pandas
conda activate covid

import numpy as np, scipy as sp, pandas as pd
from matplotlib import pyplot as plt
from sklearn.metrics.pairwise import euclidean_distances
from matplotlib import animation
from IPython.display import HTML
from itertools import combinations
from collections import namedtuple

#%load_ext line_profiler

"""
Units: time is measured in 12 hour chunks, but there's no strict correspondence to units here.
"""

class EffectiveArea:
    """Tells us the boundaries of the area that people can live in.
    """
    def __init__(self,):
        self.xmin, self.xmax = 0, 20
        self.ymin, self.ymax = 0, 20

class Person:
    def __init__(self,effectivearea,state='succeptible',distancing=False):
        self.r = 0.15
        self.days_infected = 0
        self._typical_recovery_days = 14
        self.distancing = distancing
        if self.distancing:
            self.m = 1000
        else:
            self.m = 1
        self.state = state
        self.ea = effectivearea
        self.x, self.y = np.random.uniform(self.ea.xmin,self.ea.xmax), np.random.uniform(self.ea.ymin,self.ea.ymax)
        if self.distancing:
            self.vx,self.vy = 0,0
        else:
            self.vx, self.vy = np.random.normal(size=2) # Maxwell Boltzmann?
    def move(self,dt,ea):
        self.x, self.y = self.x + self.vx*dt, self.y + self.vy*dt
        if self.state == 'infected':
            if self.days_infected < self._typical_recovery_days:
                self.days_infected += dt
            else:
                if np.random.uniform() < 0.5:
                    self.state = 'recovered'
                    self.days_infected = 0
                else:
                    self.days_infected += dt
    
def collide(p1,p2):
    if p1.state == 'infected' and p2.state == 'succeptible':
        p2.state = 'infected'
    elif p2.state == 'infected' and p1.state == 'succeptible':
        p1.state = 'infected'

    m1, m2 = p1.m, p2.m
    r1, r2 = np.array([p1.x,p1.y]), np.array([p2.x,p2.y])
    v1, v2 = np.array([p1.vx,p1.vy]), np.array([p2.vx,p2.vy])
    M = m1 + m2
    d = np.linalg.norm(r1 - r2)**2
    u1 = v1 - 2*m2 / M * np.dot(v1-v2, r1-r2) / d * (r1 - r2)
    u2 = v2 - 2*m1 / M * np.dot(v2-v1, r2-r1) / d * (r2 - r1)
    p1.vx,p1.vy = u1
    p2.vx,p2.vy = u2
    
    
    
class Universe:
    def __init__(self,npeople,initial_infection_chance=0.1,
                distancing=0.0):
        self.npeople = npeople
        self.ea = EffectiveArea()
        self.dt = 0.1
        self.data = None # gets set in self.run
        def _state():
            if np.random.uniform() < initial_infection_chance:
                return 'infected'
            return 'succeptible'
        def _distancing():
            if np.random.uniform() < distancing:
                return True
            return False
        self.people = [Person(self.ea,_state(),_distancing()) for i in range(self.npeople)]
        self.color = {'succeptible':0.5,'infected':0.0,'recovered':0.7}
        self.color = {'succeptible':'lightblue','infected':'red','recovered':'green'}
        
    def _step(self):
        points = list(zip([p.x for p in self.people],[p.y for p in self.people]))
        dists = euclidean_distances(points,points)
        close = dists < 2*self.people[0].r
        close = close.tolist()
        for (i,j) in combinations(range(self.npeople),2):
            if close[i][j]: # a bit faster than numpy indexing once things get big.
                collide(self.people[i],self.people[j])
        for p in self.people:
            p.move(self.dt,self.ea)
            if p.x < self.ea.xmin or p.x > self.ea.xmax: 
                p.vx = -p.vx
            if p.y < self.ea.ymin or p.y > self.ea.ymax:
                p.vy = -p.vy
    
    def run(self,steps):
        #x_coords[frame,particle_number]
        x_coords = np.zeros((steps,len(self.people)))
        y_coords = np.zeros((steps,len(self.people)))
        state = np.zeros((steps,len(self.people)),dtype='object')
        
        s,i,r = np.zeros(steps),np.zeros(steps),np.zeros(steps)
        def pop_count(people):
            s,i,r = 0,0,0
            for p in people:
                if p.state == 'succeptible':
                    s += 1
                elif p.state == 'infected':
                    i += 1
                elif p.state == 'recovered':
                    r += 1
            return s,i,r
        s[0],i[0],r[0] = pop_count(self.people)
        x_coords[0] = [p.x for p in self.people]
        y_coords[0] = [p.y for p in self.people]
        state[0] = [p.state for p in self.people]
        
        for step in range(1,steps):
            self._step()
            s[step],i[step],r[step] = pop_count(self.people)
            x_coords[step] = [p.x for p in self.people]
            y_coords[step] = [p.y for p in self.people]
            state[step] = [p.state for p in self.people]
        dtype = namedtuple('RunData',['x','y','state','s','i','r','steps'])
        self.data = dtype(x_coords,y_coords,state,s,i,r,steps)
    def draw(self,ax=None):
        if ax is None:
            fig,ax = plt.subplots(figsize=(5,5))
        scat = ax.scatter([p.x for p in self.people],[p.y for p in self.people],
                   c = [self.color[p.state] for p in self.people],
                   marker='.')
        return scat,

def getanim(u):
    fig,ax = plt.subplots(figsize=(4,4))

    left, width = 0.1, 0.85
    bottom, height = 0.1, 0.65
    spacing = 0.02

    rect_universe = [left, bottom, width, height]
    rect_trend = [left, bottom + height + spacing, width, 0.2]

    ax_universe = plt.axes(rect_universe)
    ax_universe.tick_params(axis='x', which='both',bottom=False,top=False,labelbottom=False)
    ax_universe.axis('off')
    ax_trend = plt.axes(rect_trend)

    s,i,r = u.data.s,u.data.i,u.data.r

    ax_trend.stackplot(range(len(s)), i, s, r, labels=['sick','healthy','recovered'],colors=['red','lightblue','green'])

    scat, = u.draw(ax_universe)
    
    def drawframe(i):
        data = np.column_stack(([u.data.x[i],u.data.y[i]]))
        scat.set_offsets(data)
        colors = np.array([u.color[_] for _ in u.data.state[i]])
        scat.set_color(colors)

        _s,_i,_r = np.zeros(len(u.data.s)),np.zeros(len(u.data.s)),np.zeros(len(u.data.s))
        _s[:i],_i[:i],_r[:i] = u.data.s[:i],u.data.i[:i],u.data.r[:i]
        #ax_trend.stackplot(range(len(_s)), _s, _i, _r, labels=['s','i','r'],colors=['blue','green','yellow'])
        #ax_trend.legend(loc='upper left')

        return scat,


    anim = animation.FuncAnimation(fig, drawframe, frames=u.data.steps,
                                  interval=20, blit=False, repeat=False)
    return anim

Now let's try it¶

For all of these, we'll do 200 people, 200 steps, initial infection chance 10%

npeople, nsteps = 200, 200
initial_infection_chance = 0.1

No distancing¶

%%capture
distancing = 0.0
u = Universe(npeople,distancing=distancing,initial_infection_chance=initial_infection_chance)
u.run(nsteps)
anim = getanim(u)

HTML(anim.to_jshtml())

25% of people distance¶

%%capture
distancing = 0.25
u = Universe(npeople,distancing=distancing,initial_infection_chance=initial_infection_chance)
u.run(nsteps)
anim = getanim(u)

HTML(anim.to_jshtml())

50% of people distance¶

%%capture
distancing = 0.5
u = Universe(npeople,distancing=distancing,initial_infection_chance=initial_infection_chance)
u.run(nsteps)
anim = getanim(u)

HTML(anim.to_jshtml())

75% of people distance¶

%%capture
distancing = 0.75
u = Universe(npeople,distancing=distancing,initial_infection_chance=initial_infection_chance)
u.run(nsteps)
anim = getanim(u)

HTML(anim.to_jshtml())

90% of people distance¶

%%capture
distancing = 0.90
u = Universe(npeople,distancing=distancing,initial_infection_chance=initial_infection_chance)
u.run(nsteps)
anim = getanim(u)